from collections.abc import Iterator
from threading import RLock
from typing import Any

import dash
import dash_bootstrap_components as dbc
from zkl_aiutils_records import DictRecordReader, RecordsScanner, get_merging_signature
from zkl_pyutils_fsspec import FsLike, resolve_fs

from .checklist import Checklist
from .graph_view import make_metrics_graph_card
from .loader import Loader
from .locking import locked


class BoardView(dash.html.Div):
    def __init__(self, fs: FsLike, style: dict | None = None, **kwargs):
        lock = RLock()
        fs = resolve_fs(fs)
        scanner = RecordsScanner(fs)

        @locked(lock)
        def load_runs_path() -> list[str]:
            runs_path = []
            for path, reader in scanner.items():
                if isinstance(reader, DictRecordReader):
                    runs_path.append(path)
            return sorted(runs_path)

        def iter_run_metrics(run_reader: DictRecordReader) -> Iterator[tuple[str, Any]]:
            for metric_name, metric_reader in run_reader.items():
                metric_signature = get_merging_signature(metric_reader)
                yield metric_name, metric_signature

        @locked(lock)
        def load_metrics() -> list[tuple[str, Any]]:
            metrics = set()
            for run_path in runs_checklist.selected_items:
                run_reader = scanner[run_path]
                assert isinstance(run_reader, DictRecordReader)
                for metric_name, metric_signature in iter_run_metrics(run_reader):
                    metrics.add((metric_name, metric_signature))
            return sorted(metrics, key=lambda x: (x[0], x[1].get('type')))

        @locked(lock)
        def load_metrics_runs_reader() -> dict[tuple[str, Any], dict[str, Any]]:
            metrics_runs_reader = {}
            for run_path in sorted(runs_checklist.selected_items):
                run_reader = scanner[run_path]
                assert isinstance(run_reader, DictRecordReader)
                run_metrics = set(iter_run_metrics(run_reader))
                run_metrics &= set(metrics_checklist.selected_items)
                for metric_name, metric_signature in sorted(run_metrics):
                    runs_reader = metrics_runs_reader.setdefault((metric_name, metric_signature), {})
                    runs_reader[run_path] = run_reader[metric_name]
            return metrics_runs_reader

        refresh_button = dbc.Button(
            children='Refresh',
            className="me-1",
            style={'flex': '0'})

        runs_path_loader = Loader(
            id='runs_reader',
            func=lambda: load_runs_path(),
            triggers=[
                dash.Input(refresh_button, 'n_clicks')])

        runs_checklist = Checklist(
            id='runs_name/checklist',
            status_triggers=[runs_path_loader.status_trigger],
            status_func=lambda: runs_path_loader.status,
            items_triggers=[runs_path_loader.value_trigger],
            items_func=lambda: runs_path_loader.value or (),
            item_label_func=lambda run_path: dash.html.Span(
                children=fs.sep + str(run_path),
                style={'padding': '0.2rem 0.3rem', 'display': 'inline'}),
            label="Runs", style={'flex': '1', 'height': '0'})

        metrics_loader = Loader(
            id='metrics',
            func=lambda: load_metrics(),
            triggers=[
                dash.Input(refresh_button, 'n_clicks'),
                runs_checklist.trigger])

        metrics_checklist = Checklist(
            id='metrics/checklist',
            status_triggers=[metrics_loader.status_trigger],
            status_func=lambda: metrics_loader.status,
            items_triggers=[metrics_loader.value_trigger],
            items_func=lambda: metrics_loader.value or (),
            item_label_func=lambda metric: dash.html.Span(
                children=[
                    dash.html.Span(metric[0]),
                    dash.html.Span(" "),
                    dash.html.Span(str(dict(metric[1])), style={'color': '#CCC'}),
                ],
                style={
                    'padding': '0.2rem 0.3rem',
                    'display': 'inline',
                }),
            label="Metrics", style={'flex': '1', 'height': '0'})

        metrics_runs_reader_loader = Loader(
            id='records',
            func=lambda: load_metrics_runs_reader(),
            triggers=[
                dash.Input(refresh_button, 'n_clicks'),
                runs_checklist.trigger, metrics_checklist.trigger])

        graphs_div = dash.html.Div(
            id='graphs',
            children=[],
            style={
                'flex': '1',
                'width': '0',
                'overflowY': 'auto',
                'display': 'flex',
                'flexDirection': 'row',
                'flexWrap': 'wrap',
            })

        @dash.callback(
            dash.Output(graphs_div, 'children'),
            metrics_runs_reader_loader.value_trigger)
        def update_graphs(*_):
            metrics_runs_reader = metrics_runs_reader_loader.value
            if metrics_runs_reader is None:
                raise dash.exceptions.PreventUpdate
            return make_metrics_graph_card(metrics_runs_reader)

        @dash.callback(
            dash.Output(refresh_button, 'disabled'),
            dash.Output(refresh_button, 'children'),
            runs_path_loader.status_trigger,
            metrics_loader.status_trigger,
            metrics_runs_reader_loader.status_trigger)
        def update_refresh_button1(*_):
            loading = any(status == 'loading' for status in (
                runs_path_loader.status,
                metrics_loader.status,
                metrics_runs_reader_loader.status))
            if loading:
                return True, 'Refreshing...'
            else:
                return False, 'Refresh'

        super().__init__(
            style={
                'display': 'flex',
                'flexDirection': 'column',
                'overflow': 'hidden',
                **(style or {}),
            },
            children=[
                runs_path_loader,
                metrics_loader,
                metrics_runs_reader_loader,

                dash.html.Div(  # header
                    style={
                        'flex': '0',
                        'padding': '0.5rem',
                        'display': 'flex',
                        'flexDirection': 'row',
                        'alignItems': 'center',
                    },
                    children=[
                        dash.html.Div(
                            children="ZKL AIUtils Board",
                            style={
                                'flex': '1',
                                'width': '0',
                                'fontWeight': 'bold'
                            }),
                        refresh_button,
                    ]),
                dash.html.Div(
                    style={
                        'flex': '0 0 1px',
                        'backgroundColor': '#e0e0e0',
                    }),
                dash.html.Div(  # body
                    style={
                        'flex': '1',
                        'height': '0',
                        'display': 'flex',
                        'flexDirection': 'row',
                    },
                    children=[
                        dash.html.Div(  # side
                            style={
                                'flex': '0 0 300px',
                                'display': 'flex',
                                'flexDirection': 'column',
                                'overflowX': 'hidden'
                            },
                            children=[
                                runs_checklist,
                                dash.html.Div(
                                    style={
                                        'flex': '0 0 1px',
                                        'backgroundColor': '#e0e0e0',
                                    }),
                                metrics_checklist,
                            ]),
                        dash.html.Div(
                            style={
                                'flex': '0 0 1px',
                                'backgroundColor': '#e0e0e0',
                            }),
                        graphs_div,
                    ])
            ],
            **kwargs)
